{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a386520",
   "metadata": {},
   "source": [
    "# 注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6a1d8b",
   "metadata": {},
   "source": [
    "本节视频为大家介绍大语言模型的核心部分注意力机制："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa6ea13f",
   "metadata": {},
   "source": [
    "![](./images/framework.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "612c1d2b",
   "metadata": {},
   "source": [
    "![](./images/seq.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a99a27a",
   "metadata": {},
   "source": [
    "## 对长句子建模过程中遇到的问题"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "050014f0",
   "metadata": {},
   "source": [
    "由于源语言与目标语言间语法结构的差异，逐词翻译文本是不可行的："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71ad90da",
   "metadata": {},
   "source": [
    "![](./images/trans.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7783bd4",
   "metadata": {},
   "source": [
    "- 在Transformer模型出现之前，​​编码器-解码器循环神经网络（Encoder-Decoder RNNs）​​是机器翻译任务中常用的架构；\n",
    "- 在此架构中，​​编码器（Encoder）​​负责处理源语言的词元（token）序列，它通过使用一种神经网络中的中间层——​​隐藏状态（hidden state）​​，来生成整个输入序列的一个​​浓缩表示（condensed representation）；"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af2158a",
   "metadata": {},
   "source": [
    "![](./images/rnn.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cbd325d",
   "metadata": {},
   "source": [
    "## 利用注意力机制捕获数据依赖性"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71752d0c",
   "metadata": {},
   "source": [
    "在网络中负责生成文本的解码器（decoder）部分，通过注意力机制能够​​选择性访问（selectively access）​​ 所有输入标记（input tokens）。这意味着在生成特定输出标记时，某些输入标记比其他标记更具重要性（hold more significance），其重要性由注意力权重决定：\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "937aa4da",
   "metadata": {},
   "source": [
    "![](./images/select.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1399bbb9",
   "metadata": {},
   "source": [
    "Transformer 中的​​自注意力（Self-attention）​​ 是一种核心技术，旨在通过允许序列中的​​每个位置（each position）​​ 与​​同一序列中的所有其他位置（every other position）​​ 进行交互并评估其​​相关性（relevance）​​，来计算出更高效、更丰富的​​输入表示（input representations）："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fd70e67",
   "metadata": {},
   "source": [
    "![](./images/attention.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943549a2",
   "metadata": {},
   "source": [
    "## ​通过自注意力机制关注输入的不同部分"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dd25d4b",
   "metadata": {},
   "source": [
    "![](./images/dot.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3ce1ebcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89], # Your     (x^1)\n",
    "   [0.55, 0.87, 0.66], # journey  (x^2)\n",
    "   [0.57, 0.85, 0.64], # starts   (x^3)\n",
    "   [0.22, 0.58, 0.33], # with     (x^4)\n",
    "   [0.77, 0.25, 0.10], # one      (x^5)\n",
    "   [0.05, 0.80, 0.55]] # step     (x^6)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "73ba6f13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = inputs[1]\n",
    "\n",
    "atte_score_2 = torch.empty(inputs.shape[0])\n",
    "\n",
    "for i, x_i in enumerate(inputs):\n",
    "    atte_score_2[i] = torch.dot(x_i, query)\n",
    "\n",
    "atte_score_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de4745ae",
   "metadata": {},
   "source": [
    "- 本小节的主要目的是演示如何以第二个输入序列$x^2$作为查询（Query），计算上下文向量$z^2$；\n",
    "- 该图描绘了此过程的初始步骤，该步骤通过点积（dot product）运算计算$x^2$与所有其他输入元素之间的注意力分数（attention scores）ω"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd1e694a",
   "metadata": {},
   "source": [
    "![](./images/dot-imp.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "155c9014",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atte_weights_2 = torch.softmax(atte_score_2, dim=0)\n",
    "\n",
    "atte_weights_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fbc78f1",
   "metadata": {},
   "source": [
    "![](./images/weights.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ec4b7a05",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.4419, 0.6515, 0.5683])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_vec_2 = torch.zeros(query.shape)\n",
    "\n",
    "for i, x_i in enumerate(inputs):\n",
    "    context_vec_2 += atte_weights_2[i] * x_i\n",
    "\n",
    "context_vec_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d3b7bf6",
   "metadata": {},
   "source": [
    "![](./images/context.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdc37b23",
   "metadata": {},
   "source": [
    "## 计算注意力权重"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cca45f2",
   "metadata": {},
   "source": [
    "推广到所有输入序列标记:\n",
    "- 上文针对特定输入的计算​​: 如前所述，我们以第二个输入序列$x^2$作为查询（query），计算了其与所有其他输入元素之间的注意力分数 ω，并据此得到了其对应的上下文向量$z^2$（如下图所示高亮行所展示的过程）；\n",
    "- 推广计算过程​​：接下来，我们将把上述计算过程​​推广​​（Generalize）到所有输入序列标记，计算​​所有​​的注意力权重和上下文向量；"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc431256",
   "metadata": {},
   "source": [
    "![](./images/all-atte.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571dbafa",
   "metadata": {},
   "source": [
    "- 在自注意力机制中，该过程始于计算注意力分数 (attention scores)，随后通过归一化处理得到总和为 1 的注意力权重 (attention weights)；\n",
    "- 这些注意力权重随后被用于通过对输入进行加权求和来生成上下文向量 (context vectors)；\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11d57f56",
   "metadata": {},
   "source": [
    "![](./images/scenario.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "02d4148d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "atte_scores = torch.empty(6, 6)\n",
    "\n",
    "for i, x_i in enumerate(inputs):\n",
    "    for j, x_j in enumerate(inputs):\n",
    "        atte_scores[i, j] = torch.dot(x_i, x_j)\n",
    "\n",
    "print(atte_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5b5e9727",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "atte_scores = inputs @ inputs.T\n",
    "print(atte_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e422762f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],\n",
      "        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],\n",
      "        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],\n",
      "        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],\n",
      "        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],\n",
      "        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])\n"
     ]
    }
   ],
   "source": [
    "atte_weights = torch.softmax(atte_scores, dim=1)\n",
    "print(atte_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "86601b3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9999)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.sum(torch.tensor([0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "25aa9473",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 3])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4849bb2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4421, 0.5931, 0.5790],\n",
       "        [0.4419, 0.6515, 0.5683],\n",
       "        [0.4431, 0.6496, 0.5671],\n",
       "        [0.4304, 0.6298, 0.5510],\n",
       "        [0.4671, 0.5910, 0.5266],\n",
       "        [0.4177, 0.6503, 0.5645]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_context_vecs = atte_weights @ inputs\n",
    "\n",
    "all_context_vecs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c052323",
   "metadata": {},
   "source": [
    "## 实现具有可训练权重的注意力模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09554d2",
   "metadata": {},
   "source": [
    "- 在本次视频中，我们将实现原始 Transformer 架构（参考论文“Attention is all you need”）、GPT 模型及大多数其他热门大语言模型（LLMs）中所使用的自注意力（self-attention）机制；\n",
    "- 该自注意力机制亦被称为“缩放点积注意力”（scaled dot-product attention）；\n",
    "    - 其整体思路与之前类似：\n",
    "        - 我们的目标仍是针对某个特定的输入元素，计算其上下文向量（context vectors），该向量是输入向量的加权和；\n",
    "        - 为此，我们仍然需要注意力权重（attention weights）；\n",
    "    - 我们将看到的，这与之前介绍的基础注意力机制仅有稍许差异：\n",
    "        - 最显著的区别是引入了在模型训练期间更新的权重矩阵（weight matrices）；\n",
    "        - 这些可训练的权重矩阵至关重要，它们使得模型（特别是模型内部的注意力模块）能够学会生成“优质”的上下文向量。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c2ac25",
   "metadata": {},
   "source": [
    "![](./images/grad-1.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "89fcf955",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 2\n"
     ]
    }
   ],
   "source": [
    "x_2 = inputs[1]\n",
    "\n",
    "d_in = inputs.shape[1]\n",
    "d_out = 2\n",
    "\n",
    "print(d_in, d_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5734e0ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(666)\n",
    "\n",
    "W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n",
    "W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n",
    "W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "588aa99e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 2])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ec6d3153",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_2.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d0fffc1",
   "metadata": {},
   "source": [
    "![](./images/grad-2.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3c03cb23",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.3927, 0.2485], grad_fn=<SqueezeBackward4>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_2 = x_2 @ W_query\n",
    "key_2 = x_2 @ W_key\n",
    "value_2 = x_2 @ W_value\n",
    "\n",
    "query_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "baf0a114",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.9699, 0.8078],\n",
       "        [1.2987, 1.4641],\n",
       "        [1.2947, 1.4453],\n",
       "        [0.6729, 0.8450],\n",
       "        [0.8586, 0.7002],\n",
       "        [0.7607, 1.0778]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keys = inputs @ W_key\n",
    "values = inputs @ W_value\n",
    "\n",
    "keys\n",
    "values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3bb22d3c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.4365, grad_fn=<DotBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_2 = keys[1]\n",
    "\n",
    "atte_score_22 = query_2.dot(key_2)\n",
    "atte_score_22"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "330a6aa6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.3092, 0.4365, 0.4303, 0.2429, 0.1968, 0.3191],\n",
       "       grad_fn=<SqueezeBackward4>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atte_score_2 = query_2 @ keys.T\n",
    "atte_score_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa355150",
   "metadata": {},
   "source": [
    "![](./images/grad-3.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "43e99b79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1648, 0.1803, 0.1795, 0.1572, 0.1522, 0.1659],\n",
       "       grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_k = keys.shape[1]\n",
    "\n",
    "atte_weights_2 = torch.softmax(atte_score_2/d_k**0.5, dim=-1)\n",
    "atte_weights_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "295ea92e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9999)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.sum(torch.tensor([0.1648, 0.1803, 0.1795, 0.1572, 0.1522, 0.1659]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8767c03b",
   "metadata": {},
   "source": [
    "![](./images/grad-4.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5005bc30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atte_weights_2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2eba851e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.9699, 0.8078],\n",
       "        [1.2987, 1.4641],\n",
       "        [1.2947, 1.4453],\n",
       "        [0.6729, 0.8450],\n",
       "        [0.8586, 0.7002],\n",
       "        [0.7607, 1.0778]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "735481c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_vec_2 = atte_weights_2 @ values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ffc7f685",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.9891, 1.0748], grad_fn=<SqueezeBackward4>)\n"
     ]
    }
   ],
   "source": [
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1202c6c",
   "metadata": {},
   "source": [
    "## 封装Self-Attention类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "623855b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f8798fb4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.9859, 1.0707],\n",
       "        [0.9891, 1.0748],\n",
       "        [0.9892, 1.0749],\n",
       "        [0.9825, 1.0657],\n",
       "        [0.9863, 1.0707],\n",
       "        [0.9825, 1.0659]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class SelfAttention_v1(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out):\n",
    "        super().__init__()\n",
    "\n",
    "        self.W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n",
    "        self.W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n",
    "        self.W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        queries = x @ self.W_query\n",
    "        keys = x @ self.W_key\n",
    "        values = x @ self.W_value\n",
    "\n",
    "        atte_scores = queries @ keys.T\n",
    "        atte_weights = torch.softmax(atte_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "        context_vecs = atte_weights @ values\n",
    "\n",
    "        return context_vecs\n",
    "    \n",
    "\n",
    "torch.manual_seed(666)\n",
    "\n",
    "sa_v1 = SelfAttention_v1(d_in, d_out)\n",
    "sa_v1(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c300331e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2125, 0.2965],\n",
       "        [0.2143, 0.2990],\n",
       "        [0.2142, 0.2989],\n",
       "        [0.2100, 0.2930],\n",
       "        [0.2099, 0.2928],\n",
       "        [0.2113, 0.2948]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class SelfAttention_v2(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, qkv_bias=False):\n",
    "        super().__init__()\n",
    "\n",
    "        self.W_query = torch.nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key = torch.nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = torch.nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "\n",
    "    def forward(self, x):\n",
    "        queries = self.W_query(x)\n",
    "        keys = self.W_key(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        atte_scores = queries @ keys.T\n",
    "        atte_weights = torch.softmax(atte_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "        context_vecs = atte_weights @ values\n",
    "\n",
    "        return context_vecs\n",
    "    \n",
    "\n",
    "torch.manual_seed(666)\n",
    "\n",
    "sa_v2 = SelfAttention_v2(d_in, d_out)\n",
    "sa_v2(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "784875f5",
   "metadata": {},
   "source": [
    "## 使用因果注意力隐藏未来词"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a8e28b",
   "metadata": {},
   "source": [
    "在因果注意力中，对角线以上的注意力权重会被掩码。这确保了对于任何给定的输入，大语言模型（LLM）在使用注意力权重计算上下文向量时，都无法利用未来的标记（future tokens）。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "351885f3",
   "metadata": {},
   "source": [
    "![](./images/mask-tri.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf05f017",
   "metadata": {},
   "source": [
    "## 因果注意力的实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5be88379",
   "metadata": {},
   "source": [
    "- 本节目标​​：我们将把先前实现的自注意力（Self-Attention）机制转换为​​因果自注意力​​（Causal Self-Attention）机制；\n",
    "- ​因果性的核心​​：因果自注意力确保了模型对序列中某个位置的预测，​​仅依赖于​​该位置之前（已知）的输出，而不会“偷看”或依赖于未来位置的信息；\n",
    "- ​​通俗理解​​：简而言之，这保证了模型在预测下一个词时，​​只能使用​​它前面的词，从而符合自回归生成的方式；\n",
    "- ​实现关键​​：为实现上述特性，对于序列中的每个当前标记（token），我们需要​​掩码（mask）掉​​输入文本中位于该标记之后的所有未来标记。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a8bd43d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.1401, 0.1888, 0.1892, 0.1548, 0.1729, 0.1541],\n",
       "        [0.1344, 0.1941, 0.1946, 0.1518, 0.1743, 0.1508],\n",
       "        [0.1348, 0.1937, 0.1943, 0.1519, 0.1743, 0.1510],\n",
       "        [0.1489, 0.1812, 0.1815, 0.1590, 0.1710, 0.1585],\n",
       "        [0.1506, 0.1800, 0.1803, 0.1591, 0.1717, 0.1582],\n",
       "        [0.1438, 0.1855, 0.1858, 0.1568, 0.1717, 0.1564]],\n",
       "       grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "queries = sa_v2.W_query(inputs)\n",
    "keys = sa_v2.W_key(inputs)\n",
    "values = sa_v2.W_value(inputs)\n",
    "\n",
    "atte_scores = queries @ keys.T\n",
    "atte_weights = torch.softmax(atte_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "atte_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ca6d7389",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 0.],\n",
       "        [1., 1., 1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_length = atte_scores.shape[0]\n",
    "\n",
    "mask_simple = torch.tril(torch.ones(context_length, context_length))\n",
    "mask_simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a630c9f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.1401, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1344, 0.1941, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1348, 0.1937, 0.1943, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1489, 0.1812, 0.1815, 0.1590, 0.0000, 0.0000],\n",
       "        [0.1506, 0.1800, 0.1803, 0.1591, 0.1717, 0.0000],\n",
       "        [0.1438, 0.1855, 0.1858, 0.1568, 0.1717, 0.1564]],\n",
       "       grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_simple_weights = atte_weights * mask_simple\n",
    "\n",
    "mask_simple_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "9976c851",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4091, 0.5909, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2578, 0.3706, 0.3716, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2220, 0.2703, 0.2706, 0.2371, 0.0000, 0.0000],\n",
       "        [0.1789, 0.2139, 0.2142, 0.1891, 0.2039, 0.0000],\n",
       "        [0.1438, 0.1855, 0.1858, 0.1568, 0.1717, 0.1564]],\n",
       "       grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "row_sums = mask_simple_weights.sum(dim=-1, keepdim=True)\n",
    "mask_simple_norm = mask_simple_weights / row_sums\n",
    "\n",
    "mask_simple_norm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bdcba36",
   "metadata": {},
   "source": [
    "![](./images/procedure.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d71a2779",
   "metadata": {},
   "source": [
    "- 从技术上讲，虽然我们现在已经完成了因果注意力（Causal Attention）机制的编码，但让我们简要了解一种更高效的方法来实现与上述相同的效果;\n",
    "- 因此，我们无需将对角线以上的注意力权重置零并对结果进行重新归一化，而是可以在​​未经归一化的注意力分数​​（unnormalized attention scores）输入 softmax 函数​​之前​​，就使用​​负无穷​​（negative infinity）来屏蔽（mask）对角线以上的区域。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8dfdcb75",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2177,   -inf,   -inf,   -inf,   -inf,   -inf],\n",
       "        [0.2706, 0.7906,   -inf,   -inf,   -inf,   -inf],\n",
       "        [0.2676, 0.7812, 0.7849,   -inf,   -inf,   -inf],\n",
       "        [0.1437, 0.4218, 0.4238, 0.2366,   -inf,   -inf],\n",
       "        [0.1395, 0.3917, 0.3940, 0.2174, 0.3245,   -inf],\n",
       "        [0.1823, 0.5425, 0.5449, 0.3053, 0.4334, 0.3017]],\n",
       "       grad_fn=<MaskedFillBackward0>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "\n",
    "masked = atte_scores.masked_fill(mask.bool(), -torch.inf)\n",
    "masked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a152c2f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4091, 0.5909, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2578, 0.3706, 0.3716, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2220, 0.2703, 0.2706, 0.2371, 0.0000, 0.0000],\n",
       "        [0.1789, 0.2139, 0.2142, 0.1891, 0.2039, 0.0000],\n",
       "        [0.1438, 0.1855, 0.1858, 0.1568, 0.1717, 0.1564]],\n",
       "       grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atte_weights = torch.softmax(masked/keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "atte_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "764a8d29",
   "metadata": {},
   "source": [
    "## 使用dropout改进因果注意力"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad450de",
   "metadata": {},
   "source": [
    "- 此外，我们还会应用​​丢弃法（Dropout）来减少训练过程中的过拟合（overfitting）；\n",
    "- Dropout可以应用于多个位置：\n",
    "    - 例如，在计算完​​注意力权重（attention weights）​​之后应用；\n",
    "    - 或者，在​​注意力权重​​与​​值向量（value vectors）​​相乘之后应用；\n",
    "- 在此，我们选择在计算完注意力权重之后应用​Dropout掩码（dropout mask），因为这种做法更为常见；\n",
    "- 另外，在此特定示例中，我们使用了50%的Dropout率（dropout rate），这意味着随机屏蔽掉一半的注意力权重（当我们后续训练GPT模型时，将使用更低的Dropout率，例如​​0.1或0.2）。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa2cf439",
   "metadata": {},
   "source": [
    "![](./images/dropout.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede5876b",
   "metadata": {},
   "source": [
    "- 若应用​0.5（即50%）的丢弃率（dropout rate）​​，未被丢弃的神经元输出值将相应地按1/0.5 = 2​的因子进行缩放；\n",
    "- 该​​缩放因子（scaling factor）的计算公式为：​​1 / (1 - dropout_rate)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "7f52c56e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "dropout = torch.nn.Dropout(0.5)\n",
    "example = torch.ones(6, 6)\n",
    "\n",
    "example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d741faeb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2., 2., 0., 2., 2., 0.],\n",
       "        [0., 0., 0., 2., 0., 2.],\n",
       "        [2., 2., 2., 2., 0., 2.],\n",
       "        [0., 2., 2., 0., 0., 2.],\n",
       "        [0., 2., 0., 2., 0., 2.],\n",
       "        [0., 2., 2., 2., 2., 0.]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dropout(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "22fabf02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5155, 0.7412, 0.7432, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.5405, 0.5413, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.4277, 0.0000, 0.3781, 0.0000, 0.0000],\n",
       "        [0.0000, 0.3710, 0.3716, 0.3137, 0.3434, 0.0000]],\n",
       "       grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "dropout(atte_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "5f5f0308",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4091, 0.5909, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2578, 0.3706, 0.3716, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2220, 0.2703, 0.2706, 0.2371, 0.0000, 0.0000],\n",
       "        [0.1789, 0.2139, 0.2142, 0.1891, 0.2039, 0.0000],\n",
       "        [0.1438, 0.1855, 0.1858, 0.1568, 0.1717, 0.1564]],\n",
       "       grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atte_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04a0adc2",
   "metadata": {},
   "source": [
    "## 因果注意力机制类的实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd5c395c",
   "metadata": {},
   "source": [
    "- 至此，我们已经准备好实现一个功能完整的自注意力（Self-Attention）机制，该实现将​​包含因果掩码（Causal Mask）和 Dropout 掩码​​；\n",
    "- 此外，还需要实现代码以支持处理​​多于一个输入​​的批次数据（Batch Data），从而确保我们的 CausalAttention 类能够兼容第1部分中数据加载器（Data Loader）所产生的​​批次输出​​；\n",
    "- 为简单起见，我们通过​​复制​​（duplicate）输入的文本样本来模拟这种批次输入；"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499dd725",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CausalAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.d_out = d_out\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape \n",
    "\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(1, 2)\n",
    "        attn_scores.masked_fill_(\n",
    "            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf) \n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "d667e466",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.4300, 0.1500, 0.8900],\n",
       "         [0.5500, 0.8700, 0.6600],\n",
       "         [0.5700, 0.8500, 0.6400],\n",
       "         [0.2200, 0.5800, 0.3300],\n",
       "         [0.7700, 0.2500, 0.1000],\n",
       "         [0.0500, 0.8000, 0.5500]],\n",
       "\n",
       "        [[0.4300, 0.1500, 0.8900],\n",
       "         [0.5500, 0.8700, 0.6600],\n",
       "         [0.5700, 0.8500, 0.6400],\n",
       "         [0.2200, 0.5800, 0.3300],\n",
       "         [0.7700, 0.2500, 0.1000],\n",
       "         [0.0500, 0.8000, 0.5500]]])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch = torch.stack((inputs, inputs), dim=0)\n",
    "\n",
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "62ccbd77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 3])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "5cf6265d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.4519,  0.2216],\n",
       "         [-0.5874,  0.0058],\n",
       "         [-0.6300, -0.0632],\n",
       "         [-0.5675, -0.0843],\n",
       "         [-0.5526, -0.0981],\n",
       "         [-0.5299, -0.1081]],\n",
       "\n",
       "        [[-0.4519,  0.2216],\n",
       "         [-0.5874,  0.0058],\n",
       "         [-0.6300, -0.0632],\n",
       "         [-0.5675, -0.0843],\n",
       "         [-0.5526, -0.0981],\n",
       "         [-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "context_length = batch.shape[1]\n",
    "dropout = 0.0\n",
    "ca = CausalAttention(d_in, d_out, context_length, dropout)\n",
    "\n",
    "context_vecs = ca(batch)\n",
    "\n",
    "context_vecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "4e2f4517",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d8595a35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0e82f6",
   "metadata": {},
   "source": [
    "![](./images/causal.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31c4f0c9",
   "metadata": {},
   "source": [
    "## 从单头注意力机制扩展至多头注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4e84f0",
   "metadata": {},
   "source": [
    "### 单个自注意力层的堆叠\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "592c9252",
   "metadata": {},
   "source": [
    "![](./images/multi-head.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd8c0577",
   "metadata": {},
   "source": [
    "![](./images/multi-procedure.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "6cf0cc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionWrapper(nn.Module):\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.heads = [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) for _ in range(num_heads)]\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([head(x) for head in self.heads], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7b6dba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "context_length = batch.shape[1]\n",
    "d_in = batch.shape[2]\n",
    "d_out = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "155a0964",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.5684,  0.5063, -0.4821,  0.4336, -0.1471,  0.4106],\n",
       "         [-0.5388,  0.6447, -0.5368,  0.5483, -0.2493,  0.3548],\n",
       "         [-0.5242,  0.6954, -0.5545,  0.5886, -0.2782,  0.3323],\n",
       "         [-0.4578,  0.6471, -0.4937,  0.5311, -0.2636,  0.2932],\n",
       "         [-0.4006,  0.5921, -0.4589,  0.5169, -0.2197,  0.2186],\n",
       "         [-0.3997,  0.5971, -0.4479,  0.4971, -0.2420,  0.2433]],\n",
       "\n",
       "        [[-0.5684,  0.5063, -0.4821,  0.4336, -0.1471,  0.4106],\n",
       "         [-0.5388,  0.6447, -0.5368,  0.5483, -0.2493,  0.3548],\n",
       "         [-0.5242,  0.6954, -0.5545,  0.5886, -0.2782,  0.3323],\n",
       "         [-0.4578,  0.6471, -0.4937,  0.5311, -0.2636,  0.2932],\n",
       "         [-0.4006,  0.5921, -0.4589,  0.5169, -0.2197,  0.2186],\n",
       "         [-0.3997,  0.5971, -0.4479,  0.4971, -0.2420,  0.2433]]],\n",
       "       grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mha = MultiHeadAttentionWrapper(d_in, d_out, context_length, 0.0, 3)\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "context_vecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "32675f76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 3])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb060d97",
   "metadata": {},
   "source": [
    "### 多头注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daf1bd34",
   "metadata": {},
   "source": [
    "- 虽然前述实现（对先前实现的单头注意力CausalAttention进行包装）是一种直观且功能完整的​​多头注意力（Multi-Head Attention）实现方式，但我们也可以编写一个名为MultiHeadAttention 的​​独立类​来达成相同目标；\n",
    "- 在此独立的MultiHeadAttention类中，我们​​并非​​通过简单拼接多个独立的单头注意力实例（CausalAttention）来构建；\n",
    "- 取而代之的方法是，我们首先创建​​单一​​的W_query、W_key 和 W_value​​权重矩阵​​，随后将这些大型矩阵的​​输出特征维度进行分割​​（split），从而为每个注意力头（attention head）生成其对应的查询、键和值表示；"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "20379683",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.1184,  0.3120, -0.0847, -0.5774],\n",
      "         [ 0.0178,  0.3221, -0.0763, -0.4225],\n",
      "         [-0.0147,  0.3259, -0.0734, -0.3721],\n",
      "         [-0.0116,  0.3138, -0.0708, -0.3624],\n",
      "         [-0.0117,  0.2973, -0.0698, -0.3543],\n",
      "         [-0.0132,  0.2990, -0.0689, -0.3490]],\n",
      "\n",
      "        [[ 0.1184,  0.3120, -0.0847, -0.5774],\n",
      "         [ 0.0178,  0.3221, -0.0763, -0.4225],\n",
      "         [-0.0147,  0.3259, -0.0734, -0.3721],\n",
      "         [-0.0116,  0.3138, -0.0708, -0.3624],\n",
      "         [-0.0117,  0.2973, -0.0698, -0.3543],\n",
      "         [-0.0132,  0.2990, -0.0689, -0.3490]]], grad_fn=<ViewBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 4])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        assert (d_out % num_heads == 0), \\\n",
    "            \"d_out must be divisible by num_heads\"\n",
    "\n",
    "        self.d_out = d_out\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = d_out // num_heads\n",
    "\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.out_proj = nn.Linear(d_out, d_out)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer(\n",
    "            \"mask\",\n",
    "            torch.triu(torch.ones(context_length, context_length),\n",
    "                       diagonal=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "\n",
    "        keys = self.W_key(x) \n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        # We implicitly split the matrix by adding a `num_heads` dimension\n",
    "        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)\n",
    "        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) \n",
    "        values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "\n",
    "        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)\n",
    "        keys = keys.transpose(1, 2)\n",
    "        queries = queries.transpose(1, 2)\n",
    "        values = values.transpose(1, 2)\n",
    "\n",
    "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
    "        attn_scores = queries @ keys.transpose(2, 3)\n",
    "\n",
    "        # Original mask truncated to the number of tokens and converted to boolean\n",
    "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
    "\n",
    "        # Use the mask to fill attention scores\n",
    "        attn_scores.masked_fill_(mask_bool, -torch.inf)\n",
    "        \n",
    "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
    "        context_vec = (attn_weights @ values).transpose(1, 2) \n",
    "        \n",
    "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
    "        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)\n",
    "        context_vec = self.out_proj(context_vec)\n",
    "\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "d_out = 4\n",
    "mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01acd775",
   "metadata": {},
   "source": [
    "- 需要注意的是，上述实现本质上是MultiHeadAttentionWrapper的一个rewritten version（重写版本），但其​​计算效率更高​​；\n",
    "- 由于​​随机权重初始化​​（random weight initializations）不同，两者产生的最终输出在数值上可能​​看起来略有差异​​，但​​两者都是功能完整的实现​​（fully functional implementations），均可用于我们后续章节中将实现的GPT类中；\n",
    "- 此外请注意，在上面的MultiHeadAttention类中，我们额外添加了一个​​线性投影层​​（linear projection layer, self.out_proj）。这只是一个不改变维度（d_out -> d_out）的线性变换。在大语言模型（LLM）的实现中，使用这样的投影层是一个​​标准惯例​​（standard convention），但它​​并非必要。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ef621ae",
   "metadata": {},
   "source": [
    "![](./images/mha.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50fd66dd",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "build-LLM-from-scratch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
